from typing import Dict, List, Any
import pandas as pd
from .ml_modules import MLModule

class WorkflowEngine:
    def __init__(self):
        self.modules: Dict[str, MLModule] = {}
        self.execution_order: List[str] = []
        self.results: Dict[str, Any] = {}

    def add_module(self, module_id: str, module: MLModule) -> None:
        """添加模块到工作流"""
        self.modules[module_id] = module

    def set_execution_order(self, order: List[str]) -> None:
        """设置模块执行顺序"""
        self.execution_order = order

    def execute(self, initial_data: pd.DataFrame = None) -> Dict[str, Any]:
        """执行工作流"""
        self.results = {}
        current_data = initial_data

        for module_id in self.execution_order:
            if module_id not in self.modules:
                raise ValueError(f"Module {module_id} not found")

            module = self.modules[module_id]
            
            # 训练模块
            if current_data is not None:
                module.fit(current_data)
            
            # 执行预测
            output_data = module.predict(current_data)
            
            # 评估结果
            evaluation = module.evaluate(output_data)
            
            # 保存结果
            self.results[module_id] = {
                "output_data": output_data,
                "evaluation": evaluation,
                "params": module.get_params()
            }
            
            # 更新当前数据
            current_data = output_data

        return self.results

    def get_module_results(self, module_id: str) -> Dict[str, Any]:
        """获取特定模块的执行结果"""
        if module_id not in self.results:
            raise ValueError(f"No results found for module {module_id}")
        return self.results[module_id]

    def visualize_results(self) -> Dict[str, Any]:
        """生成可视化结果"""
        visualization_data = {}
        for module_id, result in self.results.items():
            visualization_data[module_id] = {
                "metrics": result["evaluation"],
                "data_preview": result["output_data"].head().to_dict() if isinstance(result["output_data"], pd.DataFrame) else None
            }
        return visualization_data 